package com.mrx.xfriendship.controller;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.annotation.JSONField;
import com.mrx.xfriendship.filter.JWTFilter;
import com.mrx.xfriendship.model.dto.UserDTO;
import lombok.Builder;
import lombok.Data;
import lombok.SneakyThrows;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.context.annotation.Bean;
import org.springframework.lang.NonNull;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
import org.springframework.web.socket.handler.AbstractWebSocketHandler;
import org.springframework.web.socket.server.standard.ServletServerContainerFactoryBean;

import javax.annotation.Resource;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;

import static com.mrx.commons.config.Constant.WEBSOCKET_ATTRIBUTE_USER;

/**
 * @author : 肖康宁
 * @since : 2023-01-17 11:23
 **/
@Component
@EnableWebSocket
public class WebsocketController extends AbstractWebSocketHandler implements WebSocketConfigurer {

    private static final Logger logger = LoggerFactory.getLogger(WebsocketController.class);

    /**
     * sessionID -> session 的 map
     */
    private static final Map<String, WebSocketSession> sessions = new HashMap<>();

    /**
     * 系统用户, 发送系统消息时使用的用户
     */
    private static final UserDTO SYSTEM_USER = new UserDTO();

    @Bean
    public ServletServerContainerFactoryBean createWebSocketContainer() {
        ServletServerContainerFactoryBean container = new ServletServerContainerFactoryBean();
        // 设置 bufferSize 为 10 MB
        container.setMaxTextMessageBufferSize(10 * 1024 * 1024);
        container.setMaxBinaryMessageBufferSize(10 * 1024 * 1024);
        // 超时时间为 10 分钟
        container.setMaxSessionIdleTimeout(10 * 60 * 1000L);
        return container;
    }

    public WebsocketController() {
        SYSTEM_USER.setName("SYSTEM");
        SYSTEM_USER.setAvatar("https://p6-juejin.byteimg.com/tos-cn-i-k3u1fbpfcp/bf3fa62fcd3f4e67af7e4b180d0b6256~tplv-k3u1fbpfcp-zoom-crop-mark:3024:3024:3024:1702.awebp");
    }

    @Resource
    private JWTFilter jwtFilter;

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(this, "/chat")
                .setAllowedOrigins("*")
                .addInterceptors(jwtFilter);
    }

    @Override
    public void afterConnectionEstablished(@NonNull WebSocketSession session) {
        logger.debug("建立连接: {}-> {}", getUserNameFromSession(session), session.getId());
        sessions.put(session.getId(), session);
        sendMessage(
                SYSTEM_USER,
                session,
                new TextMessage("欢迎使用 华夏交友圈 在线聊天室, 当前在线人数为: " + sessions.size() + " 人"),
                true
        );
    }

    @Override
    public void handleTextMessage(WebSocketSession session, @NonNull TextMessage message) {
        logger.debug("收到消息: {}-> {}", session.getId(), message);
        sessions.forEach((k, to) -> {
            // 防止收到自己的消息
            if (!k.equals(session.getId())) sendMessage(session, to, message);
        });
    }

    @Override
    public void afterConnectionClosed(@NonNull WebSocketSession session, @NonNull CloseStatus status) throws Exception {
        logger.debug("关闭连接: {}: {}-> {}, code: {}",
                getUserNameFromSession(session),
                session.getId(), status.getReason(), status.getCode()
        );
        sessions.remove(session.getId()).close();
    }

    /**
     * 使用 from 中的 用户信息 发送消息给 to
     *
     * @param from    消息发送者
     * @param to      消息接收者
     * @param message 发送的消息
     * @param admin   是否为 系统消息
     */
    @SneakyThrows
    private void sendMessage(UserDTO from, WebSocketSession to, TextMessage message, boolean admin) {
        String chatMessage = ChatMessage.builder()
                .content(message.getPayload())
                .username(from.getName())
                .avatar(from.getAvatar())
                .type(ChatMessage.TYPE_RECEIVED)
                .admin(admin)
                .onlineNumber(sessions.size())
                .build()
                .buildMessage();
        logger.debug("消息转发: {}-> {}", to.getId(), chatMessage);
        to.sendMessage(new TextMessage(chatMessage));
    }

    /**
     * 使用 from 中的 用户信息 发送消息给 to
     *
     * @param from    消息发送者
     * @param to      消息接收者
     * @param message 发送的消息
     */
    private void sendMessage(WebSocketSession from, WebSocketSession to, TextMessage message) {
        sendMessage(getUserDTOFromSession(from), to, message, false);
    }


    /**
     * 从 session 中获取 JWTFilter 放进去的 UserDTO 对象, 如果没有, 则初始化一个空的 UserDTO
     *
     * @param session WebSocketSession
     * @return session 中的 UserDTO 对象
     */
    @NonNull
    private UserDTO getUserDTOFromSession(WebSocketSession session) {
        return (UserDTO) Optional.ofNullable(session.getAttributes().get(WEBSOCKET_ATTRIBUTE_USER)).orElse(new UserDTO());
    }

    @NonNull
    private String getUserNameFromSession(WebSocketSession session) {
        return getUserDTOFromSession(session).getName();
    }

    @Data
    @Builder
    private static class ChatMessage {

        /**
         * (对于客户端来说) 消息类型, 接收
         */
        @JSONField(deserialize = false)
        private static final int TYPE_RECEIVED = 0;

        private String content;

        private String username;

        private String avatar;

        private Integer type;

        private Boolean admin;

        private Integer onlineNumber;

        public String buildMessage() {
            return JSON.toJSONString(this);
        }

    }

}
